"""
    AnalysisDiagnostics

Generate GSI-compatible diagnostic output files and statistics.

This module creates diagnostic output compatible with Fortran GSI format:
- Observation space diagnostics (O-B, O-A for each observation type)
- Analysis increment fields (NetCDF output)
- Convergence diagnostics (cost function, gradient norm)
- Innovation statistics by type and region
- QC summary statistics
- Fit statistics (χ² test results)

Output files are designed to match Fortran GSI diagnostic file structure
for direct comparison and validation.
"""
module AnalysisDiagnostics

using Printf
using Statistics
using LinearAlgebra
using Dates

export DiagnosticOutput, write_innovation_diagnostics, write_analysis_increments
export write_convergence_diagnostics, write_fit_statistics, write_qc_summary
export generate_diagnostic_summary

"""
    DiagnosticOutput

Structure to hold all diagnostic output information.
"""
mutable struct DiagnosticOutput
    output_path::String
    analysis_time::DateTime

    # Observation diagnostics
    innovation_stats::Dict{String, Dict{String, Float64}}
    obs_counts::Dict{String, Dict{String, Int}}

    # Analysis diagnostics
    analysis_increments::Dict{String, Array{Float64}}
    increment_stats::Dict{String, Dict{String, Float64}}

    # Convergence diagnostics
    cost_history::Vector{Float64}
    gradient_norms::Vector{Float64}
    iterations::Int
    converged::Bool

    # QC diagnostics
    qc_summary::Dict{String, Dict{String, Int}}

    # Fit statistics
    fit_stats::Dict{String, Dict{String, Float64}}
end

"""
    DiagnosticOutput(output_path::String, analysis_time::DateTime)

Create new diagnostic output structure.
"""
function DiagnosticOutput(output_path::String, analysis_time::DateTime=now())
    return DiagnosticOutput(
        output_path,
        analysis_time,
        Dict{String, Dict{String, Float64}}(),
        Dict{String, Dict{String, Int}}(),
        Dict{String, Array{Float64}}(),
        Dict{String, Dict{String, Float64}}(),
        Float64[],
        Float64[],
        0,
        false,
        Dict{String, Dict{String, Int}}(),
        Dict{String, Dict{String, Float64}}()
    )
end

"""
    write_innovation_diagnostics(diag::DiagnosticOutput, observations, background)

Write observation space innovation diagnostics (O-B).
"""
function write_innovation_diagnostics(diag::DiagnosticOutput, observations, background)

    println("Writing innovation diagnostics...")

    filename = joinpath(diag.output_path, "innovation_diagnostics.txt")

    open(filename, "w") do f
        write(f, "GSI Julia Implementation - Innovation Diagnostics\n")
        write(f, "="^80 * "\n")
        write(f, "Analysis time: $(diag.analysis_time)\n\n")

        write(f, "Observation Type     Count     Mean O-B    RMS O-B    Bias\n")
        write(f, "-"^80 * "\n")

        for (obs_type, stats) in sort(collect(diag.innovation_stats))
            @printf(f, "%-20s %6d  %10.4f  %10.4f  %10.4f\n",
                    obs_type,
                    get(get(diag.obs_counts, obs_type, Dict()), "used", 0),
                    get(stats, "mean", 0.0),
                    get(stats, "rms", 0.0),
                    get(stats, "bias", 0.0))
        end
    end

    println("  → $filename")
end

"""
    write_analysis_increments(diag::DiagnosticOutput)

Write analysis increment fields.
"""
function write_analysis_increments(diag::DiagnosticOutput)

    println("Writing analysis increments...")

    filename = joinpath(diag.output_path, "analysis_increments.txt")

    open(filename, "w") do f
        write(f, "GSI Julia Implementation - Analysis Increments\n")
        write(f, "="^80 * "\n")
        write(f, "Analysis time: $(diag.analysis_time)\n\n")

        write(f, "Variable    Mean        Std Dev     Min         Max         RMS\n")
        write(f, "-"^80 * "\n")

        for (var, stats) in sort(collect(diag.increment_stats))
            @printf(f, "%-10s %10.4f  %10.4f  %10.4f  %10.4f  %10.4f\n",
                    var,
                    get(stats, "mean", 0.0),
                    get(stats, "std", 0.0),
                    get(stats, "min", 0.0),
                    get(stats, "max", 0.0),
                    get(stats, "rms", 0.0))
        end
    end

    println("  → $filename")
end

"""
    write_convergence_diagnostics(diag::DiagnosticOutput)

Write minimization convergence diagnostics.
"""
function write_convergence_diagnostics(diag::DiagnosticOutput)

    println("Writing convergence diagnostics...")

    filename = joinpath(diag.output_path, "convergence_diagnostics.txt")

    open(filename, "w") do f
        write(f, "GSI Julia Implementation - Convergence Diagnostics\n")
        write(f, "="^80 * "\n")
        write(f, "Analysis time: $(diag.analysis_time)\n\n")

        write(f, "Iterations: $(diag.iterations)\n")
        write(f, "Converged: $(diag.converged)\n\n")

        if !isempty(diag.cost_history)
            initial_cost = diag.cost_history[1]
            final_cost = diag.cost_history[end]
            reduction = (initial_cost - final_cost) / initial_cost * 100

            write(f, "Initial cost function: $(@sprintf("%.6e", initial_cost))\n")
            write(f, "Final cost function:   $(@sprintf("%.6e", final_cost))\n")
            write(f, "Cost reduction:        $(@sprintf("%.2f", reduction))%\n\n")
        end

        write(f, "Iteration    Cost Function    Gradient Norm\n")
        write(f, "-"^80 * "\n")

        for i in 1:length(diag.cost_history)
            @printf(f, "%9d  %16.6e  %16.6e\n",
                    i-1,
                    diag.cost_history[i],
                    i <= length(diag.gradient_norms) ? diag.gradient_norms[i] : 0.0)
        end
    end

    println("  → $filename")
end

"""
    write_fit_statistics(diag::DiagnosticOutput)

Write fit statistics (GSI fort.201-style format).
"""
function write_fit_statistics(diag::DiagnosticOutput)

    println("Writing fit statistics...")

    for (var_name, filename_num) in [("ps", 201), ("wind", 202), ("t", 203), ("q", 204)]

        filename = joinpath(diag.output_path, "fort.$filename_num")

        var_stats = get(diag.fit_stats, var_name, Dict())

        open(filename, "w") do f
            write(f, "current fit of $var_name data\n")
            write(f, "-"^80 * "\n")
            write(f, "  o-g    obs    use   type   count       bias        rms      cpen    qcpen\n")
            write(f, "-"^80 * "\n")

            count = get(var_stats, "count", 0)
            bias = get(var_stats, "bias", 0.0)
            rms = get(var_stats, "rms", 0.0)
            cpen = get(var_stats, "cpen", 0.0)
            qcpen = get(var_stats, "qcpen", 0.0)

            @printf(f, "  o-g  01   %4s  asm  %6d  %10.4f  %10.4f  %10.4f  %10.4f\n",
                    var_name, count, bias, rms, cpen, qcpen)
        end

        println("  → $filename")
    end
end

"""
    write_qc_summary(diag::DiagnosticOutput)

Write quality control summary statistics.
"""
function write_qc_summary(diag::DiagnosticOutput)

    println("Writing QC summary...")

    filename = joinpath(diag.output_path, "qc_summary.txt")

    open(filename, "w") do f
        write(f, "GSI Julia Implementation - QC Summary\n")
        write(f, "="^80 * "\n")
        write(f, "Analysis time: $(diag.analysis_time)\n\n")

        write(f, "Observation Type     Total    Used    Rejected   Reject %\n")
        write(f, "-"^80 * "\n")

        for (obs_type, qc_stats) in sort(collect(diag.qc_summary))
            total = get(qc_stats, "total", 0)
            used = get(qc_stats, "used", 0)
            rejected = get(qc_stats, "rejected", 0)
            reject_pct = total > 0 ? rejected / total * 100 : 0.0

            @printf(f, "%-20s %6d  %6d    %6d      %6.2f%%\n",
                    obs_type, total, used, rejected, reject_pct)
        end
    end

    println("  → $filename")
end

"""
    generate_diagnostic_summary(diag::DiagnosticOutput)

Generate comprehensive diagnostic summary report.
"""
function generate_diagnostic_summary(diag::DiagnosticOutput)

    println("\nGenerating diagnostic summary...")

    filename = joinpath(diag.output_path, "diagnostic_summary.txt")

    open(filename, "w") do f
        write(f, "="^80 * "\n")
        write(f, "GSI JULIA IMPLEMENTATION - DIAGNOSTIC SUMMARY\n")
        write(f, "="^80 * "\n")
        write(f, "Analysis time: $(diag.analysis_time)\n")
        write(f, "Generated: $(now())\n\n")

        # Observation summary
        write(f, "OBSERVATION SUMMARY\n")
        write(f, "-"^80 * "\n")
        total_obs = sum(get(stats, "used", 0) for stats in values(diag.obs_counts))
        write(f, "Total observations used: $total_obs\n")
        write(f, "Observation types: $(length(diag.obs_counts))\n\n")

        # Innovation summary
        write(f, "INNOVATION STATISTICS SUMMARY\n")
        write(f, "-"^80 * "\n")
        for (obs_type, stats) in sort(collect(diag.innovation_stats))
            write(f, "  $obs_type:\n")
            write(f, "    Mean O-B: $(@sprintf("%.4f", get(stats, "mean", 0.0)))\n")
            write(f, "    RMS O-B:  $(@sprintf("%.4f", get(stats, "rms", 0.0)))\n")
        end
        write(f, "\n")

        # Convergence summary
        write(f, "CONVERGENCE SUMMARY\n")
        write(f, "-"^80 * "\n")
        write(f, "Iterations: $(diag.iterations)\n")
        write(f, "Converged: $(diag.converged)\n")
        if !isempty(diag.cost_history)
            initial = diag.cost_history[1]
            final = diag.cost_history[end]
            reduction = (initial - final) / initial * 100
            write(f, "Cost reduction: $(@sprintf("%.2f", reduction))%\n")
        end
        write(f, "\n")

        # Increment summary
        write(f, "ANALYSIS INCREMENT SUMMARY\n")
        write(f, "-"^80 * "\n")
        for (var, stats) in sort(collect(diag.increment_stats))
            write(f, "  $var: RMS = $(@sprintf("%.6f", get(stats, "rms", 0.0)))\n")
        end
        write(f, "\n")

        write(f, "="^80 * "\n")
    end

    println("  → $filename")
    println("\n✓ All diagnostic files written to: $(diag.output_path)")
end

end # module AnalysisDiagnostics
